from typing import List
import tushare as ts
import pandas as pd
import matplotlib.pyplot as plt

from common.enum.form_items_key import FormItemsKey
from dto.stock_basic_dto import StockBasicDTO
from service.das.data_map_das import DataDictionaryDas

# 获取一年前的日期
dataMap = DataDictionaryDas()
result = dataMap.get_item_by_form_enum(FormItemsKey.TUSHARE_TOKEN)
ts.set_token(result.dict_value)

pro = ts.pro_api()


def getDayStock(stockCode: str, startDate: str, endDate: str):
    """
    获取股票或指数的日线数据，根据 stockCode 自动选择接口。
    :param stockCode: 股票或指数代码
    :param startDate: 开始日期 (格式：YYYYMMDD)
    :param endDate: 结束日期 (格式：YYYYMMDD)
    :return: 对应的日线数据
    """
    # 判断接口类型
    if stockCode.endswith(".SH") or stockCode.endswith(".SZ"):
        if stockCode.startswith("000") or stockCode.startswith("399"):
            # 指数代码，调用指数日线接口
            return pro.index_daily(
                ts_code=stockCode, start_date=startDate, end_date=endDate
            )
        else:
            # 股票代码，调用股票日线接口
            return pro.daily(ts_code=stockCode, start_date=startDate, end_date=endDate)
    elif (
        stockCode.endswith(".BJ")
        or stockCode.endswith(".HK")
        or stockCode.endswith(".US")
    ):
        # 北京证券交易所、港股、美股直接调用股票接口
        return pro.daily(ts_code=stockCode, start_date=startDate, end_date=endDate)
    else:
        raise ValueError(f"无法识别的股票代码格式：{stockCode}")


def getRealTime(stockCode):
    return ts.realtime_quote(ts_code=stockCode)


def get_stock_basic_info() -> List[StockBasicDTO]:
    # 获取所有股票的基础信息
    stock_basic = pro.stock_basic(
        exchange="",
        list_status="L",
        fields="ts_code,symbol,name,area,industry,cnspell,market,list_date,act_name,act_ent_type",
    )

    # 将获取的数据转换为 DTO 对象
    stock_dtos = []
    for stock in stock_basic.itertuples():
        stock_dto = StockBasicDTO(
            ts_code=stock.ts_code,
            symbol=stock.symbol,
            name=stock.name,
            area=stock.area,
            industry=stock.industry,
            cnspell=stock.cnspell,
            market=stock.market,
            list_date=stock.list_date,
            act_name=stock.act_name,
            act_ent_type=stock.act_ent_type,
        )
        stock_dtos.append(stock_dto)

    return stock_dtos
